from typing import List
from pathlib import Path
from io import BytesIO
from pdfplumber import PDF
import pdfplumber


# 定义对象
class TextObject:
    # pageNum = 0
    page_height = 0
    page_width = 0
    text = ""
    position = {}

    def __init__(self, text, page_num, position, page_height, page_width):
        self.text = text
        self.page_num = page_num
        self.position = position
        self.page_height = page_height
        self.page_width = page_width
        # px to percentage
        self.left = position['x0']/self.page_width*100
        self.top = position['top']/self.page_height*100
        self.width = position['width']/self.page_width*100
        self.height = position['height']/self.page_height*100
        # print("init position: ", self.position)
    def to_string(self):
        return 'page_num: '+str(self.page_num)+'\ntext: '+self.text+'\nposition: '+str(self.position)\
                +"left: "+str(self.left)+"top: "+str(self.top)+"width: "+str(self.width)+"height: "+str(self.height)

# 每一个TextObject都是一个单词对象
def get_word_arr(pdf_file: PDF) -> list:
    page_info = []
    for i, page in enumerate(pdf_file.pages):
        page_words = {
            "page_num": i,              # 这个是后端处理方便从0开始
            "word_object_arr": []
        }
        page_height = page.height
        page_width = page.width
        # 识别文件所有内容，获取extract_words，组装成text_object
        page_content = page.extract_words()
        for item in page_content:
            text = item['text']
            page_num = i+1              # 这个是要发给前端的从1开始
            item.pop('text')

            text_item = TextObject(text, page_num, item, page_height, page_width)
            page_words['word_object_arr'].append(text_item)
        page_info.append(page_words)
    return page_info

# 定义对象
class TextBlockObject:

    def __init__(self):
        self.text = ""
        self.size = 0
        self.text_object_arr = []
        # self.coordinates = []

    def append(self, item):
        self.text += item.text
        self.text_object_arr.append(item)
        self.size += 1

    def to_string(self):
        arr_str = ""
        for i in self.text_object_arr:
            arr_str += i.to_string()
        return 'text: ' + self.text + '\nsize: ' + str(self.size) + '\ntext_object_arr: ' + arr_str

# 定义分词器
class Tokenizer:
    chunk_size: int
    chunk_overlap: int

    def __init__(self, chunk_size, chunk_overlap):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

    def splits_page_by_words(self, pages_info: list[dict]) -> list[dict]:
        chunks_in_page = []
        for one_page in pages_info:
            chunk_list = self.splits_by_pdf_word(one_page['word_object_arr'])
            d = {
                "page_num": one_page['page_num'],
                "chunks": chunk_list
            }
            chunks_in_page.append(d)
        return chunks_in_page

    def splits_by_pdf_word(self, text_object_arr):
        block_arr = []
        block = TextBlockObject()
        appended_text_len = 0
        for itm in text_object_arr:
            if appended_text_len + len(itm.text) > self.chunk_size:
                length = self.chunk_size - appended_text_len
                half_text = TextObject(itm.text[:length], itm.page_num, itm.position, itm.page_height, itm.page_width)
                # 添加适合的文本长度
                block.append(half_text)
                # 更新已经添加文本的长度
                appended_text_len += len(half_text.text)
                # 放入结果数组中
                block_arr.append(block)
                # 新开一个chunk
                block = TextBlockObject()
                appended_text_len = 0
                # 将剩余的文本添加到新的对象中（目前这么写测试一下）
                half_text = TextObject(itm.text[length:], itm.page_num, itm.position, itm.page_height, itm.page_width)
                appended_text_len += len(half_text.text)
                block.append(half_text)
            else:
                block.append(itm)
                appended_text_len += len(itm.text)
        block_arr.append(block)
        return block_arr

def get_chunk_arr(pdf_file: PDF) -> list[dict]:
    word_arr = get_word_arr(pdf_file)
    tokenizer = Tokenizer(200, 50)
    chunks_in_page = tokenizer.splits_page_by_words(word_arr)
    return chunks_in_page

class Render:
    def __init__(self):
        # 每一个page对应一个position
        self.pages = []
        self.positions = []
        # 存储已经渲染了的页面避免重复生成-->{'1': object}
        self.images = {}

    def get_coordinates(self, chunk) -> dict:
        # 获取第一个文本片段的页码
        curr_page = chunk.text_object_arr[0].page_num
        self.pages.append(curr_page)
        pos_arr = []
        # 收集组成这一句话的所有片段的位置坐标
        for t in chunk.text_object_arr:
            # 因为同一页有好几个文本片段，所以按照页来分组讲他们归类
            # 然后按照页来渲染高亮文本框
            if curr_page == t.page_num:
                pos_arr.append(t.position)
            else:
                curr_page = t.page_num
                self.positions.append(pos_arr)
                self.pages.append(curr_page)
                pos_arr = [t.position]
        self.positions.append(pos_arr)
        return {"pages": self.pages, "coordinates": self.positions}
    # 保存指定高亮显示的文本框所在的PDF页，保存为png
    def render_page(self, pdf: PDF, text_block, file_save_path = './', file_name = 'result'):
        self.get_coordinates(text_block)
        # 然后渲染图片并保存
        for i in range(len(self.pages)):
            # 这里的pdf应该替换
            key_num = str(self.pages[i])
            if key_num not in self.images:
                img = pdf.pages[self.pages[i]].to_image()
                self.images[key_num] = img
            else:
                img = self.images[key_num]
            img.draw_rects(self.positions[i])
            name = file_name + str(self.pages[i]) + '.png'
            p1 = Path(file_save_path)
            p2 = Path(name)
            img.save(p1 / p2)
            print('Render successfully, save path is: ', p1 / p2)
        # 每次渲染完情况数组避免污染下一次渲染
        self.pages = []
        self.positions = []

def process_chunk(chunk_obj) -> dict:
    info = {
        'text': chunk_obj.text,
        'coords': []
    }
    for coord in chunk_obj.text_object_arr:
        info['coords'].append({
            'left': coord.left,
            'top': coord.top,
            'width': coord.width,
            'height': coord.height,
            'pageNum': coord.page_num
        })
    return info

def pdf_parser(pdf_pointer: BytesIO | str) -> list:
    """
    获取一个数组，里面的元素是带页码的chunks字典
    d = {'page_num': xx, 'chunks': ['text': 'xxx', 'coords': []]}
    :param pdf_pointer: 文件路径或者二进制
    :return:
    """
    with pdfplumber.open(pdf_pointer) as pdf_obj:
        chunks_in_page = get_chunk_arr(pdf_obj)
        for one_page in chunks_in_page:
            chunks = []
            for chunk_obj in one_page['chunks']:
                processed = process_chunk(chunk_obj)
                chunks.append(processed)
            one_page['chunks'] = chunks
    return chunks_in_page
